#!/usr/bin/env python3
import subprocess
import stat
import sys
import os
import os.path

def build():
    cmd = 'bash build.sh --product-name rpi3 --ccache --build-target firmware'
    subprocess.run(cmd, shell=True, check=True)

def get_block_name_by_devno(devno):
    major = os.major(devno)
    minor = os.minor(devno)
    realpath = os.readlink(F'/sys/dev/block/{major}:{minor}')
    return os.path.basename(realpath)

def ensure_devname_is_safe_to_flash(devname):
    # ensure exists
    try:
        st = os.stat(devname)
    except FileNotFoundError:
        print(F'check fail: device not found: {devname}')
        return False

    # ensure is block device
    if not stat.S_ISBLK(st.st_mode):
        print(F'check fail: not block device: {devname}')
        return False

    # get disk major:minor
    if st.st_rdev == 0:
        print(F'check fail: st_rdev is 0: {devname}')
        return False
    major = os.major(st.st_rdev)
    minor = os.minor(st.st_rdev)
    sysfs_blkdev = F'/sys/dev/block/{major}:{minor}'
    sysfs_blkdev_partition = F'{sysfs_blkdev}/partition'

    # ensure /sys/dev/block/x:y exists
    if not os.path.isdir(sysfs_blkdev):
        print(F'check fail: not a directory: {sysfs_blkdev}')
        return False

    # ensure devname not a partition
    if os.path.isfile(sysfs_blkdev_partition):
        print(F'check fail: not support partition: {devname}')
        return False

    # get all devno dependent on the block device
    devno_set = { st.st_rdev }
    diskname = get_block_name_by_devno(st.st_rdev)
    diskdir = F'/sys/class/block/{diskname}'
    for child in os.listdir(diskdir):
        childdir = os.path.join(diskdir, child)
        if not os.path.isdir(childdir):
            continue
        child_partition_file = os.path.join(childdir, 'partition')
        child_dev_file = os.path.join(childdir, 'dev')
        if not os.path.isfile(child_partition_file) or not os.path.isfile(child_dev_file):
            continue
        with open(child_dev_file, 'r') as reader:
            major, minor = reader.read().strip().split(':')
            devno_set.add(os.makedev(int(major), int(minor)))

    # ensure is not mounted
    with open('/proc/mounts', 'r') as reader:
        for line in reader:
            splited = line.strip().split()
            mount_device = splited[0]
            mount_point = splited[1]
            if mount_device and mount_device[0] == '/' and os.path.exists(mount_device):
                devno = os.stat(mount_device).st_rdev
                if devno in devno_set:
                    print(F'check fail: {mount_device} is mount on {mount_point}')
                    return False
    return True


def flash(devname):
    firmware = 'out/ohos-arm-release/packages/phone/images/firmware.img'
    if not os.path.exists(firmware):
        print(F'firmware not found: {firmware}')
        print('please run: python3 ./build-rpi3.py build')
    if not ensure_devname_is_safe_to_flash(devname):
        sys.exit(1)
    
    with open(devname, 'rb+') as writer, open(firmware, 'rb') as reader:
        filesize = reader.seek(0, 2)
        devsize = writer.seek(0, 2)
        if devsize < filesize:
            print(F'check fail: device have no enough space to flash: {filesize} < {devsize}')
            sys.exit(1)
        if devsize > 64*1024*1024*1024:
            print(F'Danger: device size more than 64GB, maybe not a sdcard: {devname}')
            print(F'If you still want to continue, run dd if={firmware} of={devname} bs=8M')
            sys.exit(1)

        reader.seek(0)
        writer.seek(0)
        bs = 8*1024*1024
        copysize = 0

        while copysize < filesize:
            readsize = bs
            leftsize = filesize - copysize
            if readsize > leftsize:
                readsize = leftsize
            data = reader.read(readsize)
            assert(len(data) == readsize)
            writesize = writer.write(data)
            assert(writesize  == readsize)
            copysize += readsize
            progress = round(copysize / filesize * 100)
            print(F'\rFlashing {devname}: {progress}%', end='')
    print(F'\nFlash success')

def help_exit(code):
    print('usgae: ')
    print('    python3 ./build-rpi3.py build')
    print('    python3 ./build-rpi3.py flash /dev/sdx')
    sys.exit(code)

if len(sys.argv) == 1:
    command = 'build'
if len(sys.argv) > 1:
    command = sys.argv[1]

if command == 'build':
    if (len(sys.argv) > 2):
        help_exit(1)
    build()
if command == 'flash':
    if (len(sys.argv) != 3):
        help_exit(1)
    flash(sys.argv[2])
